#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
#

"""
嵌套函数结构体生成器

生成函数结构体定义（异步和普通函数）
"""

import re
import sys
from mcpy.ir import IRFunctionNode
from mcpy.ir.expressions.await_expr import IRAwait
from mcpy.ir.base import IRNode
from mcpy.codegen.core.context import CodeGenContext
from mcpy.codegen.core.generator import get_global_generator
from mcpy.codegen.emitters.closure_generator import ClosureGenerator
from .async_code_generator import AsyncCodeGenerator
from .utils import get_functor_name, nested_body_generator


def add_function_struct_to_context(
    ir_module, struct_code: str, struct_name: str, prepend: bool = False
) -> None:
    """将生成的 struct 代码添加到 IR 模块的 extension 参数中

    复用 Lambda 的 lambda_structs 参数（完全兼容）

    Args:
        ir_module: IR 模块
        struct_code: struct 代码
        struct_name: struct 名称
        prepend: 如果为 True，插入到列表前面；否则追加到后面（默认）
    """
    # 复用 Lambda 的 lambda_structs 参数（嵌套函数和 Lambda 使用相同机制）
    lambda_structs = ir_module.get_extension_param("lambda_structs") or []
    if prepend:
        # 插入到前面，确保嵌套函数的结构体在父函数之前
        lambda_structs.insert(0, struct_code)
    else:
        # 追加到后面（默认行为，保持向后兼容）
        lambda_structs.append(struct_code)
    ir_module.set_extension_param(
        "lambda_structs", lambda_structs, replace_existing=True
    )


class AsyncFunctionStructGenerator:
    """异步函数结构体生成器"""

    def __init__(self):
        self.async_code_gen = AsyncCodeGenerator()

    def _generate_nested_function_structs_first(
        self, node: IRFunctionNode, context: CodeGenContext
    ) -> None:
        """递归生成函数体内所有嵌套函数的类型定义

        在生成当前异步函数结构体之前，先确保所有嵌套函数的类型已经生成。
        这样可以避免在 C++ 代码中出现前向声明问题。

        Args:
            node: 函数节点
            context: 代码生成上下文
        """
        from mcpy.codegen.emitters.nested_function_emitter import (
            NestedFunctionEmitter,
        )

        # 递归收集函数体内的所有嵌套函数
        def collect_nested_functions(stmt_node, visited=None):
            """递归收集所有嵌套函数节点"""
            if visited is None:
                visited = set()

            nested_functions = []

            # 检查当前节点是否是函数定义语句
            if hasattr(stmt_node, "node_type") and stmt_node.node_type == "function":
                # 确保不是当前函数本身
                if stmt_node is not node:
                    nested_functions.append(stmt_node)

            # 递归遍历子节点
            node_id = id(stmt_node)
            if node_id in visited:
                return nested_functions
            visited.add(node_id)

            # 从语句块中收集嵌套函数
            if hasattr(stmt_node, "statements"):
                # IRStmtBlock 或类似结构
                for stmt in stmt_node.statements:
                    nested_functions.extend(collect_nested_functions(stmt, visited))
            elif hasattr(stmt_node, "get_child_nodes"):
                # 使用 get_child_nodes 方法
                for child in stmt_node.get_child_nodes():
                    nested_functions.extend(collect_nested_functions(child, visited))
            elif hasattr(stmt_node, "__dict__"):
                # 遍历所有属性
                for attr_value in stmt_node.__dict__.values():
                    if isinstance(attr_value, IRNode):
                        nested_functions.extend(
                            collect_nested_functions(attr_value, visited)
                        )
                    elif isinstance(attr_value, list):
                        for item in attr_value:
                            if isinstance(item, IRNode):
                                nested_functions.extend(
                                    collect_nested_functions(item, visited)
                                )

            return nested_functions

        # 收集所有嵌套函数（从函数体的语句中）
        nested_functions = collect_nested_functions(node.body)

        # 获取父函数的 functor 名称，用于匹配嵌套函数
        parent_functor_name = get_functor_name(node)

        # 按照在函数体中的定义顺序排序嵌套函数（保持定义顺序）
        # 这样可以确保被依赖的函数先定义
        def get_function_source_line(func_node):
            """获取函数定义的行号"""
            return (
                func_node.source_line
                if hasattr(func_node, "source_line") and func_node.source_line > 0
                else 0
            )

        nested_functions.sort(key=get_function_source_line)

        # 为每个嵌套函数生成类型定义（递归处理，确保嵌套的嵌套函数也被处理）
        for nested_func in nested_functions:
            # 递归生成嵌套函数的结构体（如果嵌套函数内部还有嵌套函数，也会被处理）
            if nested_func.is_async:
                nested_async_gen = AsyncFunctionStructGenerator()
                nested_async_gen._generate_nested_function_structs_first(
                    nested_func, context
                )
            else:
                # 非异步嵌套函数也需要处理其内部的嵌套函数
                nested_async_gen = AsyncFunctionStructGenerator()
                nested_async_gen._generate_nested_function_structs_first(
                    nested_func, context
                )

            # 生成嵌套函数的结构体定义
            NestedFunctionEmitter.generate_struct_if_needed(nested_func, context)

            # 将嵌套函数的结构体标记为属于当前父函数
            # 这样在渲染父函数时可以找到它们
            nested_functor_name = get_functor_name(nested_func)
            ir_module = node.get_ir_module()
            if ir_module:
                # 存储嵌套函数到父函数的映射
                parent_nested_map = (
                    ir_module.get_extension_param("_parent_nested_functions") or {}
                )
                if parent_functor_name not in parent_nested_map:
                    parent_nested_map[parent_functor_name] = []
                if nested_functor_name not in parent_nested_map[parent_functor_name]:
                    parent_nested_map[parent_functor_name].append(nested_functor_name)
                ir_module.set_extension_param(
                    "_parent_nested_functions", parent_nested_map, replace_existing=True
                )

    def generate(self, node: IRFunctionNode, context: CodeGenContext) -> str:
        """生成异步函数结构体（使用 async_base 或 lambda_base）

        Args:
            node: 函数节点
            context: 代码生成上下文

        Returns:
            struct 代码
        """
        # 检查是否有 await，如果没有则使用 lambda_base（普通函数 struct）
        has_await = node.get_extension_param("has_await_expressions", False)
        if not has_await:
            function_gen = FunctionStructGenerator()
            return function_gen.generate(node, context)

        # 在生成当前异步函数结构体之前，先递归生成所有嵌套函数的类型定义
        # 这样可以确保嵌套函数的类型在使用它们的父函数结构体中已经定义
        self._generate_nested_function_structs_first(node, context)

        struct_name = get_functor_name(node)
        generator = get_global_generator()

        ret_type = node.get_cpp_return_type() or "mc::result<void>"
        m = re.match(r"mc::result<\s*(.+)\s*>", ret_type)
        ret_val_type = m.group(1) if m else "void"

        params = []
        for param in node.params:
            param_type = param.get_cpp_type() or "mc::variant"
            params.append((param.name, param_type))

        # 获取 async_dataflow（用于验证捕获变量）
        async_dataflow = node.get_extension_param("async_dataflow")

        captures = []
        captured_names = []
        captured_vars = node.captured_vars if hasattr(node, "captured_vars") else None
        if captured_vars:
            # 优先使用 captured_cpp_types 扩展参数获取准确的 C++ 类型
            captured_cpp_types = node.get_extension_param("captured_cpp_types") or {}
            for var_name in captured_vars:
                var_type = "const mc::variant&"

                # 优先使用 captured_cpp_types 中的类型
                if var_name in captured_cpp_types:
                    cpp_type = captured_cpp_types[var_name]
                    var_type = f"const {cpp_type}&"
                else:
                    # 回退：从符号表中查找类型
                    symbol = node.lookup_symbol(var_name)
                    if symbol and symbol.ir_node:
                        if isinstance(symbol.ir_node, IRFunctionNode):
                            cpp_type = symbol.ir_node.get_extension_param("cpp_type")
                            if cpp_type:
                                var_type = f"const {cpp_type}&"
                        else:
                            # 非嵌套函数变量：尝试从符号的类型信息获取
                            if symbol.type_info:
                                from mcpy.type_system import get_cpp_type

                                cpp_type = get_cpp_type(symbol.type_info, node)
                                if cpp_type:
                                    var_type = f"const {cpp_type}&"

                # 验证这确实是捕获变量（可选，用于调试）
                if async_dataflow:
                    promotion_info = async_dataflow.variable_promotions.get(var_name)
                    if promotion_info and not promotion_info.is_captured_variable:
                        import warnings

                        warnings.warn(
                            f"Variable {var_name} in captured_vars but not marked as captured variable in function {node.name}"
                        )

                captures.append((var_name, var_type))
                captured_names.append(var_name)

        if params:
            param_list = ", ".join([f"{ptype}" for _, ptype in params])
            signature = f"mc::result<{ret_val_type}>({param_list})"
        else:
            signature = f"mc::result<{ret_val_type}>()"

        # 递归收集所有 await 表达式的类型
        def collect_await_types(node_to_scan, visited=None):
            """递归遍历 IR 节点，收集所有 await 表达式的返回类型"""
            if visited is None:
                visited = set()

            # 避免循环引用
            node_id = id(node_to_scan)
            if node_id in visited:
                return []
            visited.add(node_id)

            types = []

            if isinstance(node_to_scan, IRAwait):
                await_type = (
                    node_to_scan.get_extension_param("cpp_type") or "mc::variant"
                )
                types.append(await_type)

            # 递归遍历所有子节点
            if hasattr(node_to_scan, "__dict__"):
                for attr_value in node_to_scan.__dict__.values():
                    if isinstance(attr_value, IRNode):
                        types.extend(collect_await_types(attr_value, visited))
                    elif isinstance(attr_value, list):
                        for item in attr_value:
                            if isinstance(item, IRNode):
                                types.extend(collect_await_types(item, visited))

            return types

        await_types = collect_await_types(node.body)

        if await_types:
            node.set_extension_param("await_types_list", await_types)
            unique_await_types = list(dict.fromkeys(await_types))
            node.set_extension_param("await_types_unique", unique_await_types)

        nested_context = context.create_child_context(
            current_function=node.name, current_node=node  # 设置当前节点，以便符号查找
        )
        try:
            coro_body = self.async_code_gen.generate_coro_body(
                node, ret_val_type, nested_context
            )
        except Exception as e:
            import traceback

            traceback.print_exc(file=sys.stderr)
            coro_body = """
        // 协程生成失败，使用占位实现
        throw std::runtime_error("Coroutine generation failed");"""

        # 生成两层结构，内层嵌套在外层内部
        # 内层结构体直接命名为 impl_type
        impl_name = "impl_type"

        # === 外层包装 struct ===
        outer_parts = []
        outer_parts.append(f"struct {struct_name}")
        outer_parts.append(
            f"    : public mcpy::async_base<{ret_val_type}, {struct_name}, {signature}> {{"
        )

        # === 内层实现 struct（嵌套在外层内部）===
        impl_parts = []
        impl_parts.append(f"    struct {impl_name}")
        impl_parts.append(
            f"        : public mcpy::async_impl<{ret_val_type}, {impl_name}> {{"
        )

        impl_parts.append(f"        using Self = {impl_name};")

        unique_await_types = node.get_extension_param("await_types_unique") or []
        if unique_await_types:
            types_str = ", ".join(unique_await_types)
            impl_parts.append(
                f"        using resume_value_type = mcpy::any_value_for_t<{types_str}>;"
            )
        else:
            impl_parts.append("        using resume_value_type = mcpy::any_value<>;")

        # 内层实现包含：捕获变量、参数、状态变量
        # 收集已经在 captures 中的变量名，避免在 member_variables 中重复声明
        captured_names_set = {vname for vname, _ in captures}

        for var_name, var_type in captures:
            value_type = var_type.replace("const ", "").replace("&", "").strip()
            impl_parts.append(f"        {value_type} {var_name};")

        for param_name, param_type in params:
            impl_parts.append(f"        {param_type} {param_name};")

        #   4.1 首先收集所有 await 赋值的变量（它们总是需要保存）
        #   4.2 然后收集跨 await 边界的其他变量
        async_dataflow = node.get_extension_param("async_dataflow")
        if async_dataflow:
            # 使用新的数据流分析结果
            param_names = {param[0] for param in params}

            # 定义所有成员变量（排除函数参数和函数名）
            for var_name in async_dataflow.member_variables:
                # 跳过函数参数（已在上面声明）
                if var_name in param_names:
                    continue

                # 跳过已经在 captures 中的变量（已在上面声明）
                if var_name in captured_names_set:
                    continue

                # 检查是否是嵌套函数变量
                # 如果是嵌套函数变量，需要作为成员变量生成（跨 await 使用）
                symbol = node.lookup_symbol(var_name)
                if symbol and symbol.ir_node:
                    if isinstance(symbol.ir_node, IRFunctionNode):
                        nested_func = symbol.ir_node

                        # 检查是否是简单函数（无捕获），简单函数不生成 member 变量
                        is_simple = nested_func.get_extension_param(
                            "is_simple_non_capturing", False
                        )
                        if not is_simple:
                            # 如果标记未设置，检查 captured_vars 是否为空
                            captured_vars = nested_func.captured_vars or []
                            is_simple = len(captured_vars) == 0

                        if is_simple:
                            # 简单函数不生成 member 变量，它们将在调用处通过临时对象调用
                            continue

                        # 非简单嵌套函数变量：根据存储模式生成
                        # get_functor_name 已在文件顶部导入
                        functor_name = get_functor_name(nested_func)
                        captured_vars = nested_func.captured_vars or []
                        storage_mode = nested_func.get_extension_param(
                            "storage_mode", "stack"
                        )

                        if storage_mode == "heap":
                            # 堆模式：mc::variant var_name;
                            var_type = "mc::variant"
                        else:
                            # 栈模式：functor_name var_name;
                            var_type = functor_name

                        impl_parts.append(f"    {var_type} {var_name};")
                        continue

                # 获取变量的 C++ 类型
                promotion_info = async_dataflow.variable_promotions.get(var_name)
                var_type = None

                if promotion_info:
                    from mcpy.type_system import (
                        get_cpp_type,
                        TypeKind,
                        is_protocol_type,
                    )

                    # 优先使用 type_info
                    if promotion_info.type_info:
                        var_type = get_cpp_type(promotion_info.type_info, node)

                        # 检查是否是 Protocol 类型（会映射到 mc::variant）
                        # 如果是 Protocol 类型且有初始化表达式，使用 decltype
                        if (
                            var_type == "mc::variant"
                            and is_protocol_type(promotion_info.type_info)
                            and promotion_info.init_expr
                        ):
                            # 对于 Protocol 类型，生成基于 decltype 的类型声明
                            # 使用 std::declval 避免依赖变量声明顺序
                            try:
                                from mcpy.ir.expressions.call import IRCall

                                # 检查是否是 iter() 或 next() 调用
                                if isinstance(promotion_info.init_expr, IRCall):
                                    func = promotion_info.init_expr.func
                                    func_id = getattr(func, "id", None)
                                    if (
                                        func_id == "iter"
                                        and promotion_info.init_expr.args
                                    ):
                                        # iter() 调用：生成 iterator_wrapper 类型
                                        from mcpy.type_system import TypeInfo

                                        arg_cpp_type = get_cpp_type(
                                            promotion_info.init_expr.args[
                                                0
                                            ].get_extension_param("inferred_type")
                                            or TypeInfo.any(),
                                            node,
                                        )
                                        if (
                                            arg_cpp_type
                                            and arg_cpp_type != "mc::variant"
                                        ):
                                            # 使用 std::declval 生成类型声明
                                            var_type = f"std::decay_t<decltype(mcpy::iter_wrapper(std::declval<{arg_cpp_type}>()))>"
                                        else:
                                            var_type = "mc::variant"
                                    elif func_id == "next" and promotion_info.type_info:
                                        pass
                                    else:
                                        var_type = "mc::variant"
                                else:
                                    var_type = "mc::variant"
                            except Exception:
                                # 如果生成失败，fallback 到 mc::variant
                                var_type = "mc::variant"

                    # 如果 type_info 为空，尝试从 IR 节点获取
                    if not var_type and promotion_info.ir_node:
                        inferred_type = promotion_info.ir_node.get_extension_param(
                            "inferred_type"
                        )
                        if inferred_type:
                            var_type = get_cpp_type(inferred_type, node)

                    # 如果还是为空，尝试从符号表获取
                    if not var_type and promotion_info.symbol:
                        if promotion_info.symbol.type_info:
                            var_type = get_cpp_type(
                                promotion_info.symbol.type_info, node
                            )

                # 如果仍然无法确定类型，使用 mc::variant
                if not var_type:
                    var_type = "mc::variant"

                # 转义变量名，避免与 C++ 关键字冲突
                from mcpy.codegen.configs.names import escape_cpp_name

                escaped_var_name = escape_cpp_name(var_name)
                impl_parts.append(f"        {var_type} {escaped_var_name};")

        # 内层实现的构造函数（接受所有捕获变量和参数）
        ctor_params = []
        ctor_init_list = []

        for vname, vtype in captures:
            value_type = vtype.replace("const ", "").replace("&", "").strip()
            if vtype.endswith("&"):
                ctor_params.append(f"{vtype} {vname}_")
                ctor_init_list.append(f"{vname}({vname}_)")
            else:
                ctor_params.append(f"const {value_type}& {vname}_")
                ctor_init_list.append(f"{vname}({vname}_)")

        for pname, ptype in params:
            ctor_params.append(f"{ptype} {pname}_")
            ctor_init_list.append(f"{pname}({pname}_)")

        if ctor_params:
            # 格式化构造函数参数：如果太长则换行
            params_str = ", ".join(ctor_params)
            if len(f"        {impl_name}({params_str})") > 100:
                # 参数换行
                impl_parts.append(f"        {impl_name}(")
                for i, param in enumerate(ctor_params):
                    if i < len(ctor_params) - 1:
                        impl_parts.append(f"            {param},")
                    else:
                        impl_parts.append(f"            {param})")
            else:
                impl_parts.append(f"        {impl_name}({params_str})")

            if ctor_init_list:
                impl_parts.append(f"            : {', '.join(ctor_init_list)} {{")
            else:
                impl_parts.append(f"        {{")
            impl_parts.append(f"        }}")

        # execute 方法声明（实现在源文件中）
        impl_parts.append(
            f"        mc::result<{ret_val_type}> execute(resume_value_type& __resume_value, mcpy::async_context& __ctx);"
        )

        # 内层实现结束（注意：内层结构体在外层结构体内部，所以缩进是4个空格，结束是 };）
        impl_parts.append("    };")
        impl_parts.append("")

        # 将内层结构体嵌套到外层（不需要 using 别名，因为内层结构体就叫 impl_type）

        # 外层包含捕获变量（用于传递给内层）
        for var_name, var_type in captures:
            value_type = var_type.replace("const ", "").replace("&", "").strip()
            outer_parts.append(f"        {value_type} {var_name};")

        # 外层构造函数
        if captures:
            ctor_params = []
            ctor_init_list = []
            for vname, vtype in captures:
                if vtype.endswith("&"):
                    ctor_params.append(f"{vtype} {vname}_")
                    ctor_init_list.append(f"{vname}({vname}_)")
                else:
                    ctor_params.append(f"{vtype} {vname}_")
                    ctor_init_list.append(f"{vname}(std::move({vname}_))")

            # 格式化构造函数参数：如果太长则换行
            params_str = ", ".join(ctor_params)
            if len(f"        {struct_name}({params_str})") > 100:
                # 参数换行
                outer_parts.append(f"        {struct_name}(")
                for i, param in enumerate(ctor_params):
                    if i < len(ctor_params) - 1:
                        outer_parts.append(f"            {param},")
                    else:
                        outer_parts.append(f"            {param})")
            else:
                outer_parts.append(f"        {struct_name}({params_str})")

            outer_parts.append(f"            : {', '.join(ctor_init_list)} {{")
            outer_parts.append(f"        }}")

        # operator() - 构造内层实例并调用
        if params:
            param_names = '", "'.join([pname for pname, _ in params])
            outer_parts.append(
                f"        static constexpr std::array<const char*, {len(params)}> get_param_names() {{"
            )
            outer_parts.append(f'            return {{"{param_names}"}};')
            outer_parts.append(f"        }}")

            param_decls = ", ".join([f"{ptype} {pname}" for pname, ptype in params])
            outer_parts.append(
                f"        mc::result<{ret_val_type}> operator()({param_decls}) {{"
            )

            # 构造内层实例
            impl_ctor_args = []
            for vname, _ in captures:
                impl_ctor_args.append(vname)
            for pname, _ in params:
                impl_ctor_args.append(pname)
            outer_parts.append(
                f"            return execute({', '.join(impl_ctor_args)});"
            )
            outer_parts.append(f"        }}")
        else:
            outer_parts.append(f"        mc::result<{ret_val_type}> operator()() {{")

            # 构造内层实例
            if captures:
                impl_ctor_args = [vname for vname, _ in captures]
                outer_parts.append(
                    f"            return execute({', '.join(impl_ctor_args)});"
                )
            else:
                outer_parts.append(f"            return execute();")
            outer_parts.append(f"        }}")

        # 将内层结构体添加到外层（嵌套结构）
        outer_parts.extend(impl_parts)

        # 外层结束
        outer_parts.append("};")

        # 生成完整的嵌套结构
        result = "\n".join(outer_parts)

        # 保存 execute 方法实现到扩展参数，供源文件生成使用
        # 使用 struct_name 作为键的前缀，因为 impl_name 现在是固定的 "impl_type"
        node.set_extension_param(f"{struct_name}_execute_impl", coro_body)
        node.set_extension_param(
            f"{struct_name}_impl_full_name", f"{struct_name}::impl_type"
        )
        node.set_extension_param(f"{struct_name}_ret_type", ret_val_type)

        return result

    def generate_execute_impl(
        self, node: IRFunctionNode, context: CodeGenContext
    ) -> str:
        """生成 execute 方法的实现（用于源文件）

        Args:
            node: 函数节点
            context: 代码生成上下文

        Returns:
            execute 方法的实现代码
        """
        struct_name = get_functor_name(node)

        # 获取保存的实现代码
        coro_body = node.get_extension_param(f"{struct_name}_execute_impl")
        if not coro_body:
            # 如果找不到保存的 coro_body，说明可能还没有调用过 generate()
            # 重新生成 coro_body（但不生成完整的结构体）
            ret_type = node.get_cpp_return_type() or "mc::result<void>"
            m = re.match(r"mc::result<\s*(.+)\s*>", ret_type)
            ret_val_type = m.group(1) if m else "void"

            nested_context = context.create_child_context(
                current_function=node.name, current_node=node
            )
            try:
                coro_body = self.async_code_gen.generate_coro_body(
                    node, ret_val_type, nested_context
                )
                # 保存以便后续使用
                node.set_extension_param(f"{struct_name}_execute_impl", coro_body)
                node.set_extension_param(f"{struct_name}_ret_type", ret_val_type)
                node.set_extension_param(
                    f"{struct_name}_impl_full_name", f"{struct_name}::impl_type"
                )
            except Exception as e:
                import traceback

                traceback.print_exc(file=sys.stderr)
                return ""

        # 获取返回类型
        ret_val_type = node.get_extension_param(f"{struct_name}_ret_type") or "void"

        # 获取完整的类型名称（包括外层结构体）
        full_impl_name = (
            node.get_extension_param(f"{struct_name}_impl_full_name")
            or f"{struct_name}::impl_type"
        )

        # coro_body 已经包含了完整的方法实现（包括签名和体）
        # 格式：4个空格缩进的方法签名 + 方法体（也是4个空格缩进）
        # 结构：
        #     mc::result<...> execute(...) {
        #         // header 变量
        #         co_reenter(this) {
        #             // 方法体
        #         }
        #         return set_result(__resume_value);
        #     }
        # 需要提取方法体部分（去掉签名行和第一个 {，保留方法体内的所有内容包括 co_reenter 和 return）
        lines = coro_body.split("\n")

        # 查找方法体的开始（跳过方法签名行）
        body_start_idx = 0
        for i, line in enumerate(lines):
            if "execute(resume_value_type&" in line or "execute(" in line:
                # 找到方法签名行，下一行应该是方法体开始
                body_start_idx = i + 1
                break

        # 提取方法体（从方法签名的 { 之后到方法结束的 } 之前）
        # 注意：需要保留 co_reenter(this) { ... } 和 return set_result(__resume_value);
        # coro_body 的格式：方法签名行包含 {，下一行开始就是方法体内容
        body_lines = []
        brace_count = 1  # 方法签名的 { 已经打开，所以从 1 开始

        for i in range(body_start_idx, len(lines)):
            line = lines[i]
            stripped = line.strip()

            # 计算大括号平衡
            line_brace_count = stripped.count("{") - stripped.count("}")

            # 先添加当前行
            body_lines.append(line)

            # 然后更新 brace_count
            brace_count += line_brace_count

            # 如果 brace_count 变为 0，说明方法体结束，停止提取
            if brace_count == 0:
                # 移除最后添加的行（方法体的最后一个 }）
                body_lines.pop()
                break

        adjusted_body = "\n".join(body_lines)

        # 生成方法签名和实现
        # 注意：必须使用与头文件中声明完全相同的类型
        # 头文件中使用的是 resume_value_type（类型别名），在源文件中需要使用完整的类型名称
        unique_await_types = node.get_extension_param("await_types_unique") or []
        if unique_await_types:
            types_str = ", ".join(unique_await_types)
            resume_value_type = f"mcpy::any_value_for_t<{types_str}>"
        else:
            # 如果没有 await 类型，使用默认的 any_value
            # 但需要与头文件中的声明一致
            resume_value_type = "mcpy::any_value<>"

        # 使用 resume_value_type 作为参数类型（与头文件声明一致）
        return f"mc::result<{ret_val_type}> {full_impl_name}::execute({full_impl_name}::resume_value_type& __resume_value, mcpy::async_context& __ctx) {{\n{adjusted_body}\n}}"


class FunctionStructGenerator:
    """普通函数结构体生成器"""

    def generate(self, node: IRFunctionNode, context: CodeGenContext) -> str:
        """生成嵌套函数的 C++ struct 定义（使用统一的 ClosureGenerator）

        Args:
            node: 函数节点
            context: 代码生成上下文

        Returns:
            struct 代码
        """
        generator = ClosureGenerator(
            node=node,
            context=context,
            body_fn=nested_body_generator,
            name_fn=get_functor_name,
        )
        return generator.generate()
